import torch
import copy
import torch.nn as nn
from torch.nn import functional as F
from torch.nn import Module
from torch.nn import MultiheadAttention
from torch.nn import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn import Dropout
from torch.nn import Linear
from torch.nn import LayerNorm


class Encoder(Module):
    r"""由N个编码器层堆叠而成的编码器

    Args:
        encoder_layer: EncoderLayer类的一个实例（必需）。
        num_layers: 编码器中子编码器层的数量（必需）。
        norm: 层归一化组件（可选）。
    """

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(Encoder, self).__init__()
        # 使用_get_clones函数复制encoder_layer，形成num_layers个相同的编码器层
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src):
        r"""按顺序将输入传递给编码器层

        :param src: 输入数据，通常是一个序列
        :return: 经过所有编码器层后的输出数据
        """
        output = src
        for i in range(self.num_layers):
            # 逐层传递输入
            output = self.layers[i](output)
        if self.norm:
            # 如果norm不为空，对最后一层的输出进行规范化
            output = self.norm(output)
        return output


class Decoder(Module):
    r"""由N个解码器层堆叠而成的解码器

    Args:
        decoder_layer: DecoderLayer类的一个实例（必需）。
        num_layers: 解码器中子解码器层的数量（必需）。
        norm: 层归一化组件（可选）。
    """

    def __init__(self, decoder_layer, num_layers, norm=None):
        super(Decoder, self).__init__()
        # 使用_get_clones函数复制decoder_layer，形成num_layers个相同的解码器层
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt, memory):
        r"""按照顺序将输入（和掩码）传递给解码器层

        :param tgt: 目标序列
        :param memory: 解码器在编码器阶段学到的上下文信息
        :return: 经过所有解码器层后的输出数据
        """
        output = tgt
        for i in range(self.num_layers):
            # 逐层传递输入
            output = self.layers[i](output, memory)
        if self.norm:
            # 如果norm不为空，对最后一层的输出进行规范化
            output = self.norm(output)
        return output


class EncoderLayer(Module):
    r"""编码器层，参考自CMRAN

    Args:
        d_model: 输入期望的特征数量（必需）。
        nhead: 多头注意力模型中头的数量（必需）。
        dim_feedforward: 前馈网络模型的维度（默认=1024）。
        dropout: dropout值（默认=0.1）。
        activation: 中间层的激活函数，relu或gelu（默认=relu）。
    """

    def __init__(self, d_model, nhead, dim_feedforward=1024, dropout=0.1, activation="relu"):
        super(EncoderLayer, self).__init__()
        # 初始化多头自注意力机制
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 初始化前馈网络
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)
        # 初始化归一化层
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        # 初始化dropout层
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)
        # 初始化激活函数
        self.activation = _get_activation_fn(activation)

    def forward(self, src):
        r"""将输入通过编码器层

        :param src: 输入数据
        :return: 经过自注意力和前馈网络后的输出数据
        """
        # 应用自注意力机制
        src2 = self.self_attn(src, src, src)[0]
        # 加上残差连接并进行归一化
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        # 应用前馈网络
        if hasattr(self, "activation"):
            src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        else:
            src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))
        # 再次加上残差连接并进行归一化
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src


class DecoderLayer(Module):
    r"""解码器层，参考自CMRAN

    Args:
        d_model: 输入期望的特征数量（必需）。
        nhead: 多头注意力模型中头的数量（必需）。
        dim_feedforward: 前馈网络模型的维度（默认=1024）。
        dropout: dropout值（默认=0.1）。
        activation: 中间层的激活函数，relu或gelu（默认=relu）。
    """

    def __init__(self, d_model, nhead, dim_feedforward=1024, dropout=0.1, activation="relu"):
        super(DecoderLayer, self).__init__()
        # 初始化自注意力机制
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 初始化多头注意力机制，用于编码器-解码器注意力
        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 初始化前馈网络
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)
        # 初始化归一化层
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        # 初始化dropout层
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)
        # 初始化激活函数
        self.activation = _get_activation_fn(activation)

    def forward(self, tgt, memory):
        r"""将输入（和掩码）通过解码器层

        :param tgt: 目标序列
        :param memory: 从编码器阶段传来的上下文信息
        :return: 经过多头注意力和前馈网络后的输出数据
        """
        # 将memory和tgt拼接在一起，作为多头注意力的键值对
        memory = torch.cat([memory, tgt], dim=0)
        # 应用多头注意力机制
        tgt2 = self.multihead_attn(tgt, memory, memory)[0]
        # 加上残差连接并进行归一化
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        # 应用前馈网络
        if hasattr(self, "activation"):
            tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        else:
            tgt2 = self.linear2(self.dropout(F.relu(self.linear1(tgt))))
        # 再次加上残差连接并进行归一化
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        return tgt


def _get_clones(module, N):
    r"""复制相同的模块N次，返回一个ModuleList

    :param module: 需要复制的模块
    :param N: 需要复制的次数
    :return: 包含N个相同模块的ModuleList
    """
    return ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    r"""根据字符串返回相应的激活函数

    :param activation: 激活函数的名字，字符串类型，如'relu'或'gelu'
    :return: 对应的激活函数
    """
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    else:
        raise RuntimeError("激活函数应为relu或gelu，而不是 %s." % activation)


class New_Audio_Guided_Attention(nn.Module):
    r"""新的音频引导注意力机制实现，与论文中描述的稍有不同，后发现更为高效

    Args:
        beta: 调整音频空间注意力权重的参数
    """

    def __init__(self, beta):
        super(New_Audio_Guided_Attention, self).__init__()
        # 初始化调整音频空间注意力权重的参数beta
        self.beta = beta
        self.relu = nn.ReLU()
        # 视频输入特征维度
        self.video_input_dim = 512
        # 音频输入特征维度
        self.audio_input_dim = 128
        # 隐藏层维度
        self.hidden_dim = 256
        # 通道注意力机制相关线性层
        self.affine_video_1 = nn.Linear(self.video_input_dim, self.video_input_dim)
        self.affine_audio_1 = nn.Linear(self.audio_input_dim, self.video_input_dim)
        self.affine_bottleneck = nn.Linear(self.video_input_dim, self.hidden_dim)
        self.affine_v_c_att = nn.Linear(self.hidden_dim, self.video_input_dim)
        # 空间注意力机制相关线性层
        self.affine_video_2 = nn.Linear(self.video_input_dim, self.hidden_dim)
        self.affine_audio_2 = nn.Linear(self.audio_input_dim, self.hidden_dim)
        self.affine_v_s_att = nn.Linear(self.hidden_dim, 1)

        # 潜在维度，用于更细粒度的注意力
        self.latent_dim = 4
        self.video_query = nn.Linear(self.video_input_dim, self.video_input_dim//self.latent_dim)
        self.video_key = nn.Linear(self.video_input_dim, self.video_input_dim//self.latent_dim)
        self.video_value = nn.Linear(self.video_input_dim, self.video_input_dim)
        self.affine_video_ave = nn.Linear(self.video_input_dim, self.hidden_dim)
        self.affine_video_3 = nn.Linear(self.video_input_dim, self.hidden_dim)
        self.ave_bottleneck = nn.Linear(512, 256)
        self.ave_v_att = nn.Linear(self.hidden_dim, 1)

        self.tanh = nn.Tanh()
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(0.2)
        self.norm = nn.LayerNorm(self.video_input_dim)

    def forward(self, video, audio):
        r"""音频引导的注意力机制前向传播过程

        :param video: 视频特征，形状为[batch, 10, 7, 7, 512]
        :param audio: 音频特征，形状为[batch, 10, 128]
        :return: 经过音频引导注意力机制后的视频特征，形状为[batch, 10, 512]
        """
        # 调整音频特征的时间维度顺序
        audio = audio.transpose(1, 0)
        batch, t_size, h, w, v_dim = video.size()
        a_dim = audio.size(-1)
        # 将音频特征重塑为[batch * t_size, a_dim]
        audio_feature = audio.reshape(batch * t_size, a_dim)
        # 将视频特征重塑为[batch, t_size, h*w, v_dim]
        visual_feature = video.reshape(batch, t_size, -1, v_dim)
        # 保留原始视频特征以便后续使用
        raw_visual_feature = visual_feature

        # ============================== 自注意力 =======================================
        # 将视频特征重塑为[batch * t_size, h*w, v_dim]
        video_query_feature = self.video_query(visual_feature).reshape(batch * t_size, h * w, -1)
        # 将视频特征重塑为[batch * t_size, v_dim, h*w]
        video_key_feature = self.video_key(visual_feature).reshape(batch * t_size, h * w, -1).permute(0, 2, 1)
        # 计算注意力能量矩阵
        energy = torch.bmm(video_query_feature, video_key_feature)
        # 应用softmax函数得到注意力权重矩阵
        attention = self.softmax(energy)
        # 将视频特征重塑为[batch * t_size, h*w, v_dim]
        video_value_feature = self.video_value(visual_feature).reshape(batch * t_size, h * w, -1)
        # 将注意力权重矩阵与视频特征矩阵相乘得到加权后的视频特征
        output = torch.matmul(attention, video_value_feature)
        # 对加权后的视频特征加上dropout并进行归一化
        output = self.norm(visual_feature.reshape(batch * t_size, h * w, -1) + self.dropout(output))
        # 将加权后的视频特征重塑为[batch, t_size, h*w, v_dim]
        visual_feature = output.reshape(batch, t_size, h*w, v_dim)

        # ============================== 视频空间注意力 ================================
        # 计算视频特征的平均值，作为全局上下文
        video_average = visual_feature.sum(dim=1)/(h*w)
        # 将平均值重塑为[batch*t_size, v_dim]
        video_average = video_average.reshape(batch*t_size, v_dim)
        # 将全局上下文通过线性层并应用relu激活函数
        video_average = self.relu(self.affine_video_ave(video_average)).unsqueeze(-2)
        # 将视频特征重塑为[batch*t_size, -1, v_dim]
        self_video_att_feat = visual_feature.reshape(batch*t_size, -1, v_dim)
        # 计算视频特征的查询矩阵，并应用relu激活函数
        self_video_att_query = self.relu(self.affine_video_3(self_video_att_feat))
        # 将全局上下文与查询矩阵相乘
        self_query = self_video_att_query * video_average
        # 计算空间注意力权重矩阵
        self_spatial_att_maps = self.softmax(self.tanh(self.ave_v_att(self_query)).transpose(2,1))
        # 将空间注意力权重矩阵与视频特征相乘得到最终的加权视频特征
        self_att_feat = torch.bmm(self_spatial_att_maps, visual_feature).squeeze().reshape(batch, t_size, v_dim)

        # ============================== 通道注意力 ==================================
        # 计算音频特征的查询矩阵，并应用relu激活函数
        audio_query_1 = self.relu(self.affine_audio_1(audio_feature)).unsqueeze(-2)
        # 计算视频特征的查询矩阵，并应用relu激活函数
        video_query_1 = self.relu(self.affine_video_1(visual_feature)).reshape(batch*t_size, h*w, -1)
        # 将音频特征与视频特征的查询矩阵相乘，计算通道注意力
        audio_video_query_raw = (audio_query_1 * video_query_1).mean(-2)
        # 将通道注意力通过线性层并应用relu激活函数
        audio_video_query = self.relu(self.affine_bottleneck(audio_video_query_raw))
        # 计算通道注意力矩阵，并应用sigmoid函数
        channel_att_maps = self.affine_v_c_att(audio_video_query).sigmoid().reshape(batch, t_size, -1, v_dim)
        # 将通道注意力矩阵与原始视频特征相乘得到通道加权后的视频特征
        c_att_visual_feat = (raw_visual_feature * (channel_att_maps + 1))

        # ============================== 音频空间注意力 ==============================
        # 将通道加权后的视频特征重塑为[batch * 10, 49, v_dim]
        c_att_visual_feat = c_att_visual_feat.reshape(batch*t_size, -1, v_dim)
        # 计算通道加权后的视频特征的查询矩阵，并应用relu激活函数
        c_att_visual_query = self.relu(self.affine_video_2(c_att_visual_feat))
        # 计算音频特征的查询矩阵，并应用relu激活函数
        audio_query_2 = self.relu(self.affine_audio_2(audio_feature)).unsqueeze(-2)
        # 将通道加权后的视频特征的查询矩阵与音频特征的查询矩阵相乘
        audio_video_query_2 = c_att_visual_query * audio_query_2
        # 计算空间注意力权重矩阵
        spatial_att_maps = self.softmax(self.tanh(self.affine_v_s_att(audio_video_query_2)).transpose(2, 1))
        # 将空间注意力权重矩阵与通道加权后的视频特征相乘得到最终的加权视频特征
        c_s_att_visual_feat = torch.bmm(spatial_att_maps, c_att_visual_feat).squeeze().reshape(batch, t_size, v_dim)

        # ============================== 合并注意力结果 ==============================
        # 将自注意力的结果和音频空间注意力的结果结合，并输出最终结果
        c_s_att_visual_feat = c_s_att_visual_feat + self.beta*self_att_feat.sigmoid()*c_s_att_visual_feat
        return c_s_att_visual_feat
